import gradio as gr
import numpy as np
import pickle
from PIL import Image

# 加载KNN模型
with open('best_knn_model.pkl', 'rb') as f:
    knn_model = pickle.load(f)

def predict_digit(input_data):
    print("Input type:", type(input_data))
    print("Input structure:", input_data)

    if input_data is None:
        return "请在画板上绘制一个数字"
    
    try:
        # 从输入字典中提取图像数据
        if isinstance(input_data, dict) and 'composite' in input_data:
            image = input_data['composite']
        else:
            return "无效的输入数据"

        # 确保输入是 NumPy 数组
        image = np.array(image)
        
        print("Image shape after conversion:", image.shape)
        print("Image dtype:", image.dtype)

        # 转换为灰度图像
        if image.ndim == 3:
            image = np.mean(image[:,:,:3], axis=2)  # 只取 RGB 通道，忽略 Alpha 通道
        
        # 对图像进行二值化处理
        threshold = 128
        image = (image > threshold).astype(np.uint8) * 255
        
        # 将输入的手写图像转换为 PIL 格式的灰度图，并调整大小为8x8
        image = Image.fromarray(image.astype('uint8')).convert('L').resize((8, 8))
        
        # 将图像数据转换为数组，并将其展平为1D数组
        img_array = np.array(image).reshape(1, -1)
        
        # 反转颜色：手写数字通常是黑色背景上的白色
        img_array = 255 - img_array
        
        # 进行预测
        prediction = knn_model.predict(img_array)
        return f"预测结果: {int(prediction[0])}"
    except Exception as e:
        print(f"Error: {str(e)}")
        return f"处理图像时出错: {str(e)}"

# 创建Gradio界面
with gr.Blocks(theme=gr.themes.Base()) as iface:
    gr.Markdown("# 手写数字识别")
    
    with gr.Row():
        with gr.Column(scale=1):
            # 使用Sketchpad替代Image组件
            image_input = gr.Sketchpad(height=280, width=280)
        
        with gr.Column(scale=1):
            output = gr.Textbox(label="输出")
            flag_button = gr.Button("标记")
    
    submit_button = gr.Button("提交")
    clear_button = gr.Button("清除")

    # 设置按钮点击事件
    submit_button.click(predict_digit, inputs=image_input, outputs=output)
    clear_button.click(lambda: None, inputs=None, outputs=image_input)

if __name__ == "__main__":
    iface.launch(share=True)